package com.leetcode.No0927;

import java.util.Arrays;

/**
 * @program: Solution
 * @description: 三等分
 * @author: Wang Zhihua
 * @date: 2022-10-08
 */
public class Solution {

    public int[] threeEqualParts(int[] arr) {
        int oneNum = 0; // 三组01数组中1的数量
        for (int num : arr) {
            oneNum += num;
        }

        if (oneNum % 3 != 0) {
            return new int[]{-1, -1};
        }
        if (oneNum == 0) {
            return new int[]{0, 2};
        }

        oneNum /= 3;
        int len = arr.length;
        int leftX = -1, rightX = -1, leftY = -1, rightY = -1;
        int tempOneNum = 0;
        for (int i = 0; i < len; ++i) {
            tempOneNum += arr[i];
            if (tempOneNum == oneNum) {
                if (leftX < 0) {
                    leftX = i;
                }
            }
            if (tempOneNum == oneNum + 1) {
                if (rightX < 0) {
                    rightX = i;
                }
            }
            if (tempOneNum == (oneNum << 1)) {
                if (leftY < 0) {
                    leftY = i;
                }
            }
            if (tempOneNum == (oneNum << 1) + 1) {
                if (rightY < 0) {
                    rightY = i;
                }
            }
        }

        int leftZeroNum = 0;
        for (int i = 0; i < len - 1; ++i) {
            if (arr[i] == 1) {
                break;
            }
            ++leftZeroNum;
        }

        int rightZeroNum = 0;
        for (int i = len - 1; i >= 0; --i) {
            if (arr[i] == 1) {
                break;
            }
            ++rightZeroNum;
        }

        if ((rightX - leftX - 1 < rightZeroNum) || (rightY - leftY - 1 < rightZeroNum)) {
            return new int[]{-1, -1};
        }

        int x = Math.max(leftX - leftZeroNum + 1, leftY - rightX + 1);
        x = Math.max(x, len - rightY - rightZeroNum);
        for (int i = 0; i < x; ++i) {
            if (arr[leftZeroNum + i] != arr[rightX + i] || arr[rightX + i] != arr[rightY + i]) {
                return new int[]{-1, -1};
            }
        }

        return new int[]{leftX + rightZeroNum, leftY + rightZeroNum + 1};
    }

    public static void main(String[] args) {
        Solution solution = new Solution();

        int[] arr = new int[]{1,0,1,0,1};
        System.out.println(Arrays.toString(solution.threeEqualParts(arr)));

        arr = new int[]{1,1,0,1,1};
        System.out.println(Arrays.toString(solution.threeEqualParts(arr)));

        arr = new int[]{1,1,0,0,1};
        System.out.println(Arrays.toString(solution.threeEqualParts(arr)));

        arr = new int[]{
                1,1,1,0,1,0,0,1,0,1,0,0,0,1,0,0,
                1,1,1,0,1,0,0,1,0,1,0,0,0,1,0,0,
                0,0,1,1,1,0,1,0,0,1,0,1,0,0,0,1,0,0};
        System.out.println(Arrays.toString(solution.threeEqualParts(arr)));
    }

}
